import numpy as np
import matplotlib.pyplot as plt
from data_loader_2D import load_dataset
from rknn_dpc import RKNN_DPC
from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score
import time
import matplotlib
matplotlib.use('TkAgg')  # 切换到 TkAgg 后端
matplotlib.rcParams['font.sans-serif'] = ['SimHei']
matplotlib.rcParams['axes.unicode_minus'] = False

def evaluate_clustering(true_labels, pred_labels):
    # 去除未分配的点
    mask = pred_labels != -1
    true_filtered = true_labels[mask]
    pred_filtered = pred_labels[mask]
    
    if len(np.unique(pred_filtered)) < 2:
        return 0, 0
    
    ari = adjusted_rand_score(true_filtered, pred_filtered)
    nmi = normalized_mutual_info_score(true_filtered, pred_filtered)
    
    return ari, nmi

# 参数搜索
def parameter_search(dataset_name='Spiral'):
    """
    RKNN-DPC: K近邻数取[5, 51)，步长为1，预备值取[2, K]，步长为1
    """

    # 加载数据
    X, true_labels = load_dataset(dataset_name)
    print(f"\n=== {dataset_name} 数据集的参数搜索 ===")
    print(f"数据形状: {X.shape}")
    print(f"真实类别数: {len(np.unique(true_labels))}")
    
    # 设置参数范围
    k_values = range(5, 51)  # [5, 51)，步长为1
    
    # 存储结果
    results = []
    best_result = {'ari': -1, 'k': 0, 'R': 0}
    
    # 获取真实类簇数
    true_m = len(np.unique(true_labels))
    
    print(f"\n开始搜索参数组合...")
    print(f"K范围: {k_values.start} 到 {k_values.stop-1}")
    print(f"对于每个K，R范围: 2 到 K")
    
    start_time = time.time()
    total_combinations = sum(k-1 for k in k_values)  # R从2到K
    current_combination = 0
    
    for k in k_values:
        print(f"\nK的取值为：{k}")
        r_values = range(2, k+1)  # [2, K]，步长为1
        
        for R in r_values:
            current_combination += 1
            progress = (current_combination / total_combinations) * 100 # 打印搜索进度，方便查看
            
            try:
                # 运行聚类
                rknn_dpc = RKNN_DPC(k=k, m=true_m, R=R)
                rknn_dpc.rknn_dpc(X)
                labels = rknn_dpc.labels
                
                # 评估结果
                n_clusters = len(np.unique(labels[labels != -1]))
                ari, nmi = evaluate_clustering(true_labels, labels)
                
                result = {
                    'k': k,
                    'R': R,
                    'n_clusters': n_clusters,
                    'ari': ari,
                    'nmi': nmi
                }
                results.append(result)
                
                # 更新最佳结果
                if ari > best_result['ari']:
                    best_result = result.copy()
                    print(f"  新最佳结果: K={k}, R={R}, ARI={ari:.3f}, 进度={progress:.1f}%")
                elif ari > 0.9:  # 打印高质量结果
                    print(f"  高质量结果: K={k}, R={R}, ARI={ari:.3f}, 进度={progress:.1f}%")
                else:
                    pass
                    # print(f"  当前K值结果不佳, 进度={progress:.1f}%")
                
            except Exception as e:
                print(f"  K={k}, R={R} 失败: {str(e)[:50]}")
    
    elapsed_time = time.time() - start_time
    print(f"\n搜参完成，总耗时: {elapsed_time:.1f}秒")
    print(f"总共搜索了 {len(results)} 个参数组合")
    
    # 分析结果
    if results:
        # 按ARI排序，找到前10个最佳结果
        results.sort(key=lambda x: x['ari'], reverse=True)
        top_results = results[:10]
        
        print(f"\n=== 前10个最佳参数的组合 ===")
        print(f"{'排名':>4} {'K':>4} {'R':>4} {'类簇数':>8} {'ARI':>8} {'NMI':>8}")
        print("-" * 50)
        
        for i, result in enumerate(top_results, 1):
            print(f"{i:>4} {result['k']:>4} {result['R']:>4} "
                  f"{result['n_clusters']:>8} {result['ari']:>8.3f} {result['nmi']:>8.3f}")
        
        # 可视化最佳结果
        best = top_results[0]
        print(f"\n最终最佳参数: K={best['k']}, R={best['R']}, ARI={best['ari']:.3f}")
        
        # 最佳参数的可视化
        rknn_dpc = RKNN_DPC(k=best['k'], m=true_m, R=best['R'])
        rknn_dpc.rknn_dpc(X)
        labels = rknn_dpc.labels
        
        # 可视化（二维数据集）
        if X.shape[1] == 2:
            plt.figure(figsize=(15, 5))
            
            # 真实标签
            plt.subplot(1, 2, 1)
            plt.scatter(X[:, 0], X[:, 1], c=true_labels, cmap='tab10', s=30, alpha=0.8)
            plt.title(f'{dataset_name} - 真实标签')
            plt.xlabel('X1')
            plt.ylabel('X2')
            
            # 最佳聚类结果
            plt.subplot(1, 2, 2)
            unique_labels = np.unique(labels)
            colors = plt.cm.tab10(np.linspace(0, 1, max(3, len(unique_labels))))
            
            for i, label in enumerate(unique_labels):
                if label != -1:
                    plt.scatter(X[labels == label, 0], X[labels == label, 1], 
                               c=[colors[label]], s=30, alpha=0.8, label=f'类簇{label}')
            
            # 标记类簇中心
            plt.scatter(X[rknn_dpc.centers, 0], X[rknn_dpc.centers, 1], 
                       s=200, edgecolors='red', facecolors='none', 
                       linewidths=2, marker='*', label='类簇中心')
            
            plt.title(f'{dataset_name} - 最佳聚类结果\n(K={best["k"]}, R={best["R"]}, ARI={best["ari"]:.3f})')
            plt.xlabel('X1')
            plt.ylabel('X2')
            plt.legend()

            # plt.savefig(f'{dataset_name}_parameter_search.png', dpi=150)
            plt.show()
        
        return top_results
    else:
        print("没有找到有效结果")
        return []

def test_best_known():
    """测试最佳参数"""
    print("\n=== 测试最佳参数 ===")
    
    # 已知的最佳参数
    best_params = {
        'Jain': {'k': 15, 'R': 5},
    }
    
    for dataset_name, params in best_params.items():
        try:
            X, true_labels = load_dataset(dataset_name)
            true_m = len(np.unique(true_labels))
            
            rknn_dpc = RKNN_DPC(k=params['k'], m=true_m, R=params['R'])
            rknn_dpc.rknn_dpc(X)
            labels = rknn_dpc.labels
            
            ari, nmi = evaluate_clustering(true_labels, labels)
            n_clusters = len(np.unique(labels[labels != -1]))
            
            print(f"{dataset_name}: K={params['k']}, R={params['R']} -> "
                  f"类簇数={n_clusters}, ARI={ari:.3f}, NMI={nmi:.3f}")
            
        except Exception as e:
            print(f"{dataset_name} 测试失败: {e}")

if __name__ == "__main__":  

    # 只测试已知最佳参数
    # test_best_known()

    # 指定数据集的全面搜索
    dataset = "Jain"
    parameter_search(dataset)